# 创建输入dataloader，定义损失优化器等等，使用模型fit函数进行训练
import torch
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import torch.optim as optim
from train import fit
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
from Config import config  # 从自定义Config.py参数文件中插入
from Model import Model
import os
from sklearn.metrics import confusion_matrix
import seaborn as sns

torch.cuda.empty_cache()
# 1、加载数据
config = config()

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 使用GPU进行训练
print(device)

# 定义数据处理
transformation = transforms.Compose([    transforms.ToTensor()  # ToTensor（）作用： 1、转化一个tensor；2、将数值归一化，即转换到0-1之间；3、将channel维度放在第一维上
])

# 2、创建数据集dataloader
# datasets.MNIST是pytorch中提供的一个类，用于加载MNIST手写数字数据集
# config.batch_size指定了数据集的存储路径，如果该路径下没有数据集，会根据download=True参数自动下载
# train=True 表面我们正在加载【训练数据集】，MNIST数据集分为训练集和测试集，train=True加载的是训练集
# transform=transformation 应用在每个样本上转换，transformation就是上面自定义的张量设置，维度设置，将PIL图像转换为张量，归一化等操作
train_ds = datasets.MNIST(config.data_path,train=True,download=True,transform=transformation)
# datasets.MNIST是pytorch中提供的一个类，用于加载MNIST手写数字数据集
# config.batch_size指定了数据集的存储路径，如果该路径下没有数据集，会根据download=True参数自动下载
# train=False 表面我们正在加载【测试数据集】，MNIST数据集分为训练集和测试集，train=False加载的是测试集
# transform=transformation 应用在每个样本上转换，transformation就是上面自定义的张量设置，维度设置，将PIL图像转换为张量，归一化等操作
test_ds = datasets.MNIST(config.data_path,train=False,download=True,transform=transformation)
#
train_dl = DataLoader(train_ds,batch_size=config.batch_size,shuffle=True)
test_dl = DataLoader(test_ds,batch_size=config.batch_size,shuffle=False)

# 观察每个输入的样本形状
imgs, labels = next(iter(train_dl))
print("每个输入的样本形状为{}".format(imgs.shape))
# torch.size([61,1,28,28])  【批次、channel、高、宽】

# 3、导入模型
model = Model().to(device)

# 4、创建损失函数
loss_fn = torch.nn.CrossEntropyLoss()
# 5、创建优化器
optim = optim.Adam(model.parameters(),lr=config.learning_rate)

# 6、初始化所诉列表和文件·
train_loss = []
train_acc = []
test_loss = []
test_acc = []
bst_loss = np.inf

df1 = pd.DataFrame(columns=['train_loss','test_loss'])
df1.to_csv('./loss.csv',index=False)
df2 = pd.DataFrame(columns=['train_acc','test_acc'])
df2.to_csv('./acc.csv',index=False)  # 路径可以根据需要更改

# 7、开始训练并存储数据到csv
for epoch in range(config.epochs):
    epoch_loss,epoch_acc,test_epoch_loss,test_epoch_acc = fit(epoch,model,train_dl,test_dl,loss_fn,optim,device,bst_loss)
    list1 = [epoch_loss,test_epoch_loss]  # 创建存放loss的列表
    data1 = pd.DataFrame([list1])
    data1.to_csv('./loss.csv',mode='a',header=False,index=False)

    list2 = [epoch_acc,test_epoch_acc] # 创建存放acc的列表
    data2 = pd.DataFrame([list2])
    data2.to_csv('./acc.csv', mode='a', header=False, index=False)

    train_loss.append(epoch_loss)
    train_acc.append(epoch_acc)
    test_loss.append(test_epoch_loss)
    test_acc.append(test_epoch_acc)

# 8、训练结果可视化
plt.plot(range(1, config.epochs+1), train_loss, label='train_loss')
plt.plot(range(1, config.epochs+1), test_loss, label='test_loss')
plt.legend()
plt.show()
plt.plot(range(1, config.epochs+1), train_acc, label='train_acc')
plt.plot(range(1, config.epochs+1), test_acc, label='test_acc')
plt.legend()
plt.show()
